package com.zidongxiangxi.sharding.job.core.executor;

import com.zidongxiangxi.sharding.job.core.executor.runnable.ShardingJobRunnable;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.FutureTask;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 线程池分片任务执行者
 *
 * @author chenxudong
 * @date 2021/02/18
 */
@Slf4j
public class ThreadPoolShardingJobExecutor implements ShardingJobExecutor {
    private static final int MAX_TERMINATE_RETRY_TIMES = 5;

    private final ThreadPoolExecutor threadPool = new ThreadPoolExecutor(
        Integer.MAX_VALUE, Integer.MAX_VALUE, 10, TimeUnit.SECONDS,
        new SynchronousQueue<>(),
        new ThreadFactoryBuilder().setNameFormat("SHARDING-ITEM-THREAD-%d").build(),
        new ThreadPoolExecutor.AbortPolicy());
    private final Set<String> shardingItems = new TreeSet<>();
    private final Map<String, FutureTask> shardingItemTaskMap = new ConcurrentHashMap<>();
    private final String jobName;
    private final ShardingJobRunnable runnable;

    private volatile boolean shutdown = false;

    public ThreadPoolShardingJobExecutor(String jobName, ShardingJobRunnable runnable) {
        this.threadPool.allowCoreThreadTimeOut(true);
        this.jobName = jobName;
        this.runnable = runnable;
    }

    @Override
    public synchronized void execute(WaitingExecuteShardingItem waitingExecuteShardingItem) {
        if (shutdown) {
            log.warn("任务【{}】执行器已经关闭", jobName);
            return;
        }
        String shardingItem = waitingExecuteShardingItem.getShardingItem();
        if (shardingItems.contains(waitingExecuteShardingItem.getShardingItem())) {
            log.warn("任务【{}】的分片【{}】已经存在", jobName, shardingItem);
            return;
        }
        FutureTask futureTask = (FutureTask)threadPool.submit(() -> {
            waitingExecuteShardingItem.start();
            while (!Thread.interrupted()) {
                try {
                    runnable.run(shardingItem);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            }
            log.warn("任务：{}， 分片：{}，被中断，结束分片任务的执行", jobName, shardingItem);
            waitingExecuteShardingItem.stop();
        });
        shardingItems.add(shardingItem);
        shardingItemTaskMap.put(shardingItem, futureTask);
    }

    @Override
    public synchronized void finishShardingItem(String shardingItem) {
        if (shutdown) {
            return;
        }
        shardingItems.remove(shardingItem);
        FutureTask futureTask = shardingItemTaskMap.remove(shardingItem);
        if (Objects.nonNull(futureTask)) {
            futureTask.cancel(true);
        }
    }

    @Override
    public void finishAllShardingItems() {
        List<FutureTask> futureTasks = new ArrayList<>(shardingItemTaskMap.values());
        shardingItems.clear();
        shardingItemTaskMap.clear();
        for (FutureTask futureTask : futureTasks) {
            futureTask.cancel(true);
        }
    }

    @Override
    public synchronized void shutdown() {
        shutdown = true;
        shardingItems.clear();
        shardingItemTaskMap.clear();
        threadPool.shutdownNow();
        int retryTimes = 0;
        while (retryTimes < MAX_TERMINATE_RETRY_TIMES) {
            boolean interrupted = false;
            boolean termination = false;
            try {
                termination = threadPool.awaitTermination(2, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                interrupted = true;
            }
            if (interrupted || termination) {
                break;
            }
            retryTimes++;
        }
    }

}
